from typing import Optional, Type
from asyncer import asyncify
from lagent.actions.base_action import AsyncActionMixin, BaseAction, tool_api
from lagent.actions.parser import BaseParser, JsonParser
from lagent.schema import ActionReturn, ActionStatusCode


class ArxivSearch(BaseAction):
    """
    arXiv 论文搜索工具。
    用于搜索物理学、数学、计算机科学、生物学、金融、统计学、
    电气工程和经济学等领域的科学论文。
    """

    def __init__(
        self,
        top_k_results: int = 3,  # 返回的最大结果数
        max_query_len: int = 300,  # 最大查询长度
        doc_content_chars_max: int = 1500,  # 文档内容最大字符数
        description: Optional[dict] = None,  # 工具描述
        parser: Type[BaseParser] = JsonParser,  # 结果解析器
    ):
        super().__init__(description, parser)
        self.top_k_results = top_k_results
        self.max_query_len = max_query_len
        self.doc_content_chars_max = doc_content_chars_max

    @tool_api(explode_return=True)
    def get_arxiv_article_information(self, query: str) -> dict:
        """
        执行 arXiv 搜索并获取文章元信息。

        Args:
            query: 搜索查询内容

        Returns:
            dict: 文章信息
                * content (str): 包含3篇论文信息的列表
        """
        import arxiv

        try:
            # 执行搜索查询
            results = arxiv.Search(
                query[: self.max_query_len],  # 限制查询长度
                max_results=self.top_k_results  # 限制返回结果数
            ).results()
        except Exception as exc:
            # 处理搜索异常
            return ActionReturn(
                errmsg=f'Arxiv exception: {exc}', 
                state=ActionStatusCode.HTTP_ERROR
            )

        # 格式化每篇论文的信息
        docs = [
            f'Published: {result.updated.date()}\n'  # 发布日期
            f'Title: {result.title}\n'  # 论文标题
            f'Authors: {", ".join(a.name for a in result.authors)}\n'  # 作者列表
            f'Summary: {result.summary[:self.doc_content_chars_max]}'  # 摘要(限制长度)
            for result in results
        ]

        if docs:
            # 返回合并后的论文信息
            return {'content': '\n\n'.join(docs)}
        # 没有找到结果时的返回
        return {'content': 'No good Arxiv Result was found'}


class AsyncArxivSearch(AsyncActionMixin, ArxivSearch):
    """
    异步版本的 arXiv 搜索工具。
    提供与同步版本相同的功能，但支持异步调用。
    """

    @tool_api(explode_return=True)
    @asyncify  # 将同步方法转换为异步方法
    def get_arxiv_article_information(self, query: str) -> dict:
        """
        异步执行 arXiv 搜索并获取文章元信息。

        Args:
            query: 搜索查询内容

        Returns:
            dict: 文章信息
                * content (str): 包含3篇论文信息的列表
        """
        # 调用父类的同步方法
        return super().get_arxiv_article_information(query)
